import argparse
import logging
import math
import os
import sys
import time
import pandas
from consts import DB_NAME, LONG, SHORT, BUY, SELL
from leach import Leach
from order import Order
from position import Position
from pymongo import MongoClient
from signal import Signal

logging.basicConfig(
    stream=sys.stdout, level=logging.INFO, format='%(asctime)s-%(levelname)s-%(message)s'
)


class Worker(object):
    intervalMap = {
        "H": "h", "D": "d", "W": "d", "M": "d"
    }

    def __init__(self, directory, url):
        self.directory = directory
        self.client = MongoClient(url)
        self.leach = Leach(directory)
        self.order = Order()
        self.position = Position()
        self.signal = Signal()

    @property
    def db(self):
        return self.client[DB_NAME]

    @staticmethod
    def is_close(direction, side):
        if direction == LONG and side == SELL:
            return True
        if direction == SHORT and side == BUY:
            return True
        return False

    def fit_signals(self, dataframe, models, market):
        dfs = list()
        for ids in models.values():
            models = list()
            for model in self.db.model.find({"_id": {"$in": ids}}):
                models.append(model)
            temp_df = self.leach.result(dataframe, models, market)
            dfs.append(temp_df)
        if not dfs:
            return pandas.DataFrame()
        dfs = pandas.concat(dfs)
        dfs = dfs.drop_duplicates("symbol")
        return dfs.reset_index(drop=True)

    def get_signals(self, dataframe, document):
        dfs = list()
        for direction in document["config"]:
            for side in document["config"][direction]:
                signals = self.fit_signals(dataframe, document["config"][direction][side]["models"], document["market"])
                signals.loc[signals.index, ["direction", "side"]] = [direction, side]
                dfs.append(signals)
        if not dfs:
            return pandas.DataFrame()
        dfs = pandas.concat(dfs)
        return dfs.reset_index(drop=True)

    def get_position(self, symbol, direction):
        temp_df = self.position.query(symbol=dict(eq=symbol), direction=dict(eq=direction))
        if temp_df.shape[0]:
            return temp_df.iloc[0]
        else:
            return None

    def open_deviate(self, tick, signal, params):
        period = self.intervalMap[signal["interval"]]
        deviate = (tick[f"{period}_open"] - signal[f"{period}_close"]) / signal[f"{period}_close"] * 100
        # logging.info(f'{tick["symbol"]}开盘价格偏离为{round(deviate, 2)}%')
        min_open_deviate = params.get(f"{signal['direction']}_min_open_deviate")
        if min_open_deviate is not None and deviate < min_open_deviate:
            return False
        max_open_deviate = params.get(f"{signal['direction']}_max_open_deviate")
        if max_open_deviate is not None and deviate > max_open_deviate:
            return False
        return True

    def stop_by_drawdown(self, position, params, tick):
        drawdown_rate = params.get('drawdown_rate')
        if drawdown_rate:
            period = self.intervalMap[position.interval]
            if position.direction == LONG:
                stop_price = position.upper_price * (1 - drawdown_rate / 100)
                price = tick[f"{period}_low"]
                if price < stop_price:
                    return price
            if position.direction == SHORT:
                stop_price = position.floor_price * (1 + drawdown_rate / 100)
                price = tick[f"{period}_high"]
                if price > stop_price:
                    return price

    def stop_by_expired(self, position, params, tick):
        position_expired_days = params.get('position_expired_days')
        if position_expired_days is None:
            return False
        if position.step_count > position_expired_days:
            period = self.intervalMap[position.interval]
            return tick[f"{period}_open"]

    def stop_by_price(self, position, params, tick):
        stop_price_rate = params.get('stop_price_rate')
        if stop_price_rate is not None:
            period = self.intervalMap[position.interval]
            if position.direction == LONG:
                stop_price = position.stop_price - position.stop_price * (stop_price_rate / 100)
                price = tick[f"{period}_low"]
                if price <= stop_price:
                    return stop_price
            else:
                stop_price = position.stop_price + position.stop_price * (stop_price_rate / 100)
                price = tick[f"{period}_high"]
                if price >= stop_price:
                    return stop_price

    def stop_by_signal(self, position, tick):
        if position.direction == LONG:
            signals = self.signal.query(
                symbol=dict(eq=position.symbol),
                side=dict(eq=SELL), direction=dict(eq=LONG), state=dict(eq=0)
            )
        else:
            signals = self.signal.query(
                symbol=dict(eq=position.symbol),
                side=dict(eq=BUY), direction=dict(eq=SHORT), state=dict(eq=0)
            )
        self.signal.dataframe[signals.index, "state"] = 1
        if signals.shape[0] > 0:
            period = self.intervalMap[position.interval]
            return tick[f"{period}_open"]

    def stop_by_target(self, position, params, tick):
        stop_rate = params.get('stop_rate')
        if stop_rate:
            price = position.cost_price
            period = self.intervalMap[position.interval]
            target_rate = params.get('target_rate')
            slope_rate = params.get('slope_rate') or 0
            if position.direction == LONG:
                count = 0
                if target_rate:
                    upper = position.upper_price
                    count = (upper - price) / price * 100 / target_rate
                    count = math.floor(count) if count > 0 else 0
                stop_price = price * (1 - stop_rate / 100 + slope_rate / 100 * count)
                if tick[f"{period}_low"] <= stop_price:
                    return stop_price
            if position.direction == SHORT:
                count = 0
                if target_rate:
                    floor = position.floor_price
                    count = (floor - price) / price * 100 / target_rate
                    count = math.ceil(count) if count < 0 else 0
                stop_price = price * (1 + stop_rate / 100 + slope_rate / 100 * count)
                if tick[f"{period}_high"] >= stop_price:
                    return stop_price

    @staticmethod
    def lock_position(position, params):
        lock_position_days = params.get('lock_position_days') or 0
        return position.step_count < lock_position_days

    def lock_profit(self, position, params, tick):
        lock_profit_rate = params.get('lock_profit_rate')
        if not lock_profit_rate:
            return False
        period = self.intervalMap[position.interval]
        profit = tick[f"{period}_open"] - position.cost_price
        return profit / abs(position.cost_price_price) * 100 < lock_profit_rate

    def close_position(self, position, datetime, price, msg):
        profit_rate = (price - position.cost_price) / position.cost_price
        if position.direction == SHORT:
            profit_rate = (position.cost_price - price) / position.cost_price
        self.order.insert(
            symbol=position.symbol, direction=position.direction,
            open_datetime=position.create_datetime, open_price=position.cost_price,
            close_datetime=datetime, close_price=price, profit_rate=profit_rate,
            step_count=position.step_count, msg=msg
        )
        self.position.delete(symbol=dict(eq=position.symbol), direction=dict(eq=position.direction))

    def open_position(self, tick, signal):
        period = self.intervalMap[signal["interval"]]
        stop_price = signal[f"{period}_low"] if signal["direction"] == LONG else signal[f"{period}_high"]
        self.position.insert(
            symbol=tick["symbol"], direction=signal["direction"], interval=signal["interval"],
            create_datetime=tick["datetime"], update_datetime=tick["datetime"], cost_price=tick[f"{period}_open"], stop_price=stop_price,
            upper_price=tick[f"{period}_open"], floor_price=tick[f"{period}_open"], step_count=0
        )

    def insert_order(self):
        orders = self.order.query()
        self.db.deal.insert_many(orders.to_dict(orient="records"))

    @staticmethod
    def get_tick(ticks, symbol):
        temp_df = ticks.loc[ticks["symbol"].eq(symbol)]
        if temp_df.shape[0]:
            return temp_df.iloc[0].to_dict()
        else:
            return None

    def settle_positions(self, ticks):
        positions = self.position.query()
        for _, position in positions.iterrows():
            row_data = position.to_dict()
            tick = self.get_tick(ticks, row_data["symbol"])
            period = self.intervalMap[row_data["interval"]]
            row_data.update({
                "step_count": row_data["step_count"] + 1,
                "update_datetime": tick["datetime"]
            })
            if tick[f"{period}_high"] > row_data["upper_price"]:
                row_data["upper_price"] = tick[f"{period}_high"]
            if tick[f"{period}_low"] < row_data["floor_price"]:
                row_data["floor_price"] = tick[f"{period}_low"]
            self.position.update(row_data, symbol=dict(eq=row_data["symbol"]), direction=dict(eq=row_data["direction"]))

    def close_signal(self, signal):
        positions = self.position.query(symbol=dict(eq=signal["symbol"]), direction=dict(eq=signal["direction"]))
        if positions.shape[0]:
            return self.signal.insert(**signal, state=0, msg="信号平仓")
        return None

    def open_signal(self, signal):
        position = self.get_position(signal["symbol"], signal["direction"])
        if position is None:
            return self.signal.insert(**signal, state=0, msg="信号开仓")
        return None

    def close_unsettled_position(self, ticks):
        for _, position in self.position.dataframe.iterrows():
            tick = self.get_tick(ticks, position.symbol)
            if tick is None:
                continue
            period = self.intervalMap[position.interval]
            self.close_position(position, tick["datetime"], tick[f"{period}_close"], "末日平仓")

    def run_close(self, ticks, params):
        positions = self.position.query()
        for _, position in positions.iterrows():
            tick = self.get_tick(ticks, position.symbol)
            if tick is None:
                continue
            if self.lock_position(position, params):
                continue
            if self.lock_profit(position, params, tick):
                continue
            stop_price = self.stop_by_expired(position, params, tick)
            if stop_price is not None:
                self.close_position(position, tick["datetime"], stop_price, "到期平仓")
                continue
            stop_price = self.stop_by_signal(position, tick)
            if stop_price is not None:
                self.close_position(position, tick["datetime"], stop_price, "拐点平仓")
                continue
            stop_price = self.stop_by_drawdown(position, params, tick)
            if stop_price is not None:
                self.close_position(position, tick["datetime"], stop_price, "回撤止盈")
                continue
            stop_price = self.stop_by_price(position, params, tick)
            if stop_price is not None:
                self.close_position(position, tick["datetime"], stop_price, "价格止损")
                continue
            stop_price = self.stop_by_target(position, params, tick)
            if stop_price is not None:
                self.close_position(position, tick["datetime"], stop_price, "移动止损")
                continue

    def run_open(self, ticks, params):
        signals = self.signal.query(state=dict(eq=0))
        for _, row in signals.iterrows():
            row_data = row.to_dict()
            tick = ticks.loc[ticks["symbol"].eq(row_data["symbol"])].iloc[0].to_dict()
            if self.open_deviate(tick, row_data, params):
                self.open_position(tick, row_data)
        self.signal.dataframe.loc[signals.index, "state"] = 1

    def run_task(self, document):
        temp_df = pandas.DataFrame()
        rootpath = f"{self.directory}/{document['market']}/snapshot"
        for filename in sorted(os.listdir(rootpath)):
            filepath = f"{rootpath}/{filename}"
            logging.info(filepath)
            if os.path.isfile(filepath) and os.path.splitext(filename)[1] == ".csv":
                temp_df = pandas.read_csv(filepath, low_memory=False)
                self.settle_positions(temp_df)
                signals = self.get_signals(temp_df, document)
                for _, signal in signals.iterrows():
                    row_data = signal.to_dict()
                    row_data["interval"] = document["interval"]
                    if self.is_close(row_data["direction"], row_data["side"]):
                        self.close_signal(row_data)
                    else:
                        self.open_signal(row_data)
                self.run_close(temp_df, document["params"])
                self.run_open(temp_df, document["params"])
        if temp_df.shape[0] and self.position.dataframe.shape[0]:
            self.close_unsettled_position(temp_df)

    def get_res(self):
        temp_df = self.order.query()
        l_temp_df = temp_df.loc[temp_df["direction"].eq(LONG)]
        s_temp_df = temp_df.loc[temp_df["direction"].eq(SHORT)]
        result = {
            "total_symbols": len(set(temp_df["symbol"].values.tolist())),
            "total_signals": temp_df.shape[0],
            "total_profit_rate": temp_df["profit_rate"].sum(),
            "total_wins": temp_df.loc[temp_df["profit_rate"].gt(0)].shape[0],
            "total_loss": temp_df.loc[temp_df["profit_rate"].lt(0)].shape[0],
            "min_profit_rate": temp_df["profit_rate"].min(),
            "max_profit_rate": temp_df["profit_rate"].max(),
            "long_signals": l_temp_df.shape[0],
            "long_profit_rate": l_temp_df["profit_rate"].sum(),
            "long_wins": l_temp_df.loc[temp_df["profit_rate"].gt(0)].shape[0],
            "long_loss": l_temp_df.loc[temp_df["profit_rate"].lt(0)].shape[0],
            "long_min_profit_rate": l_temp_df["profit_rate"].min(),
            "long_max_profit_rate": l_temp_df["profit_rate"].max(),
            "short_signals": s_temp_df.shape[0],
            "short_profit_rate": s_temp_df["profit_rate"].sum(),
            "short_wins": s_temp_df.loc[temp_df["profit_rate"].gt(0)].shape[0],
            "short_loss": s_temp_df.loc[temp_df["profit_rate"].lt(0)].shape[0],
            "short_min_profit_rate": s_temp_df["profit_rate"].min(),
            "short_max_profit_rate": s_temp_df["profit_rate"].max()
        }
        result.update({
            "total_win_rate": result["total_wins"] / result["total_loss"],
            "long_win_rate": result["long_wins"] / result["long_loss"],
            "short_win_rate": result["short_wins"] / result["short_loss"]
        })
        return result

    def start(self):
        for document in self.db.backtest.find({"state": 0}):
            try:
                logging.info(document)
                self.db.backtest.update_one(
                    filter={"_id": document["_id"]}, update={
                        "$set": {
                            "state": 1
                        }
                    })
                self.order = Order()
                self.position = Position()
                strategy = self.db.strategy.find_one({"_id": document["strategy"]})
                self.run_task(strategy)
                res = self.get_res()
            except Exception as exp:
                logging.error(exp)
                self.db.backtest.update_one(
                    filter={"_id": document["_id"]}, update={
                        "$set": {
                            "state": 3, "error": str(exp)
                        }
                    })
            else:
                self.db.backtest.update_one(
                    filter={"_id": document["_id"]}, update={
                        "$set": {
                            "state": 2, **res
                        }
                    })
        self.client.close()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--directory", type=str, default="/pdfs")
    parser.add_argument("--url", type=str, default="mongodb://admin:8VvvQ8SpHvctIR2G@106.15.54.243:27017/admin?authSource=admin")
    params, _ = parser.parse_known_args()
    while True:
        Worker(params.directory, params.url).start()
        time.sleep(3)


if __name__ == "__main__":
    main()
